""" 检测模块 """
import base64
import queue
import threading
import time

from PIL import Image

from io import BytesIO
from datetime import datetime

from yolo import YOLO_ONNX


class DetectParking:
    def __init__(self):
        # 存储base64编码的图像信息
        self.queue = queue.Queue()
        # 每隔10秒向外界获取一次最新的图像信息
        self.timer = 10
        # 开启获取图片的线程
        self.start_get_picture()
        self.yolo = self.model_load()


    def get_img(self):
        """ python队列, 无数据会进行阻塞 """
        return self.queue.get()

    def get_picture(self, sleep_time):
        """ 获取图片转为base64编码 """
        while True:
            # todo 获取图片
            with open(r"C:/Users/ruiyangz/Desktop/yolov5-pytorch-main/img/demo.jpg", "rb") as f:
                data = f.read()
                msg = base64.b64encode(data)
                self.queue.put(msg)
                print("获取图片成功", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            time.sleep(sleep_time)

    def start_get_picture(self):
        task = threading.Thread(target=self.get_picture, args=(self.timer,))
        task.start()

    def base64_to_img(self, img_base64):
        # 解码Base64数据
        image_data = base64.b64decode(img_base64)

        # 使用 BytesIO 将二进制数据转换为文件对象
        image_stream = BytesIO(image_data)

        # 使用 plt 读取图像数据
        image = Image.open(image_stream)
        return image

    def detect(self):
        img_base64 = self.get_img()
        img = self.base64_to_img(img_base64)
        print("detecting...")
        res = self.yolo.detect_image(img)
        print(f"detect finished... result is {res}")
        return res

    def model_load(self):
        """ 加载模型 """
        print("loading model...")
        res = YOLO_ONNX()
        print("model loaded successful...")
        return res


if __name__ == '__main__':
    detect = DetectParking()
    detect.detect()
